"""
光学目标识别模块，增强机载相机拍摄画面并初步识别目标
此模块基于 ultralytics 的 YOLO 框架实现

date: 2025/9/21
author: SiHeng Tang
file: optics_scanner.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import logging
import time
from copy import deepcopy
from typing import Any

from async_sche import TaskWrapper, BaseTaskGen, ExitGenerationError
from helpers.history_interpolation import DistantHistoryError
from mengji import Vehicle
from mod_loader import ConfigHandler
from task_element import DataEndpoint

logger = logging.getLogger(__name__)
handler = ConfigHandler(__name__)

DEFAULT_CONFIG = {
    "label": "optics_scanner",
    "flight_history": "optics_fp",
    "result_cache_cnt": 2,
    "min_update_interval": 1,
    "enable_yolo_postprocess": 0,
    "yolo_model_path": "/home/gicg/ec130/tools/cam/model"
}


class OpticsScanner(BaseTaskGen):
    """
    光学扫描模块，分析 camera 节点产生的位置图像数据，拆分目标图像，打包目标图像和原始图像
    图像包::

    地图总览 (overview):
        {
            "overview": base64 ndarray,
            "shot_time": int,
            "position":
            {
                "lat": float latitude,
                "lon": float longitude
            }
        }
    目标寻址 (target_view):
        {
            "target_view": base64 ndarray
        }
    """

    def __init__(self, config: dict, veh: Vehicle):
        super().__init__(priority=TaskWrapper.IDLE_PRIORITY + 10)
        self._veh = veh
        self._fp_node_name = config["flight_history"]
        self._result_cache_max = config["result_cache_cnt"]
        self._min_tx_interval = config["min_update_interval"]
        self._model_path = config["yolo_model_path"]
        # self._overview_trans_fmt: tuple[int, str] = (config["overview_quality"], config["overview_type"])
        # self._target_view_trans_fmt: tuple[int, str] = (config["target_view_quality"], config["target_view_type"])

        self.recv_frame = DataEndpoint()
        self._fp = None

        self._active_telemetry = False
        self._last_tx_time = time.time()
        self._scanner_overview: list[dict] = []
        self._target_result: dict[int, Any] = {}

        logger.debug(f"loading yolo model at {self._model_path}")
        # self._model = YOLO(self._model_path)

    def _on_gen(self):
        if not self.recv_frame.detect_new_in_dp():
            raise ExitGenerationError()

    async def _func(self):
        frame_pld = self.recv_frame.get_dp().get_a("payload")

        if not bool(frame_pld):
            return

        # frame_content = frame_pld["frame"]
        frame_timestamp = frame_pld["cap_time"]

        # 飞行历史检索，获得拍摄位置
        # 如果获取不到历史记录就没有必要进行下一步，因为没有位置信息都是白搭
        if self._fp is None:
            self._fp = self._veh.get_a(self._fp_node_name, None)
            return

        try:
            _gps = self._fp.get_history_at(time.time() - frame_timestamp).get_a("payload")
        except DistantHistoryError:
            logger.warning(f"distant history error when scanning target, please check system load")
            return

        # 目标检测
        # result = self._model(frame_content)

        # 结果封装
        if self._active_telemetry:
            res = {
                "overview": frame_pld["str_frame"],
                "shot_time": frame_timestamp,
                "position": {
                    "lat": _gps["lat"],
                    "lon": _gps["lon"],
                    "hdg": _gps["hdg"],
                    "alt": _gps["relative_alt"]
                }
            }
            if len(self._scanner_overview) >= self._result_cache_max:
                self._scanner_overview = self._scanner_overview[1:]
            self._scanner_overview.append(res)

    def _on_exit(self):
        self.recv_frame = None

    def command_QUERY_TARGET_VIEW_BY_ID(self, rpc_params: int):
        """
        获取目标 ID 对应的图像
        :param rpc_params: 整数目标 ID
        :return: 目标图像存在则返回有效数据包，不存在返回 0
        """
        if not isinstance(rpc_params, int):
            return 0

        if rpc_params in self._target_result:
            return self._target_result.get(rpc_params, 0)
        else:
            return 0

    def topic_SCANNER_OVERVIEW(self, rpc_params: str):
        """
        目标区域总览话题
        :param rpc_params: 话题参数
        :return: 最近的一次数据存在则返回有效数据包，不存在返回 0
        """
        if rpc_params == "subscribe":
            self._active_telemetry = True
            return 0

        if rpc_params == "unsubscribe":
            self._active_telemetry = False
            self._scanner_overview.clear()
            return 0

        if rpc_params == "mailbox_update":
            if len(self._scanner_overview) > 0:
                if (time.time() - self._last_tx_time) < self._min_tx_interval:
                    return 0

                self._last_tx_time = time.time()
                res = self._scanner_overview[0]
                self._scanner_overview = self._scanner_overview[1:]
                return res
            else:
                return 0

        return 0


@handler.add_handler("unit_proc:optics_scanner")
def load_config(veh: Vehicle, cfg_section: dict):
    logger.debug(f"optics scanner load config {cfg_section}")

    cfg = deepcopy(DEFAULT_CONFIG)
    cfg.update(cfg_section)

    veh.reg_part(cfg["label"], node := OpticsScanner(cfg, veh), is_sche=True)
    veh.get_a("api").add_topic("scanner/overview", node.topic_SCANNER_OVERVIEW)
    veh.get_a("api").add_command("scanner/target_view", node.command_QUERY_TARGET_VIEW_BY_ID)
